import torch
import torch.nn as nn

# 示例参数
vocab_size = 5  # 词汇表大小为5
batch_size = 2  # 批次大小为2
seq_length = 3  # 序列长度为3

# 创建 BigramLanguageModel 实例
class BigramLanguageModel(nn.Module):
    def __init__(self, vocab_size):
        super().__init__()
        # each token directly reads off the logits for the next token from a lookup table
        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)

    def forward(self, idx, targets=None):
        # idx and targets are both (B,T) tensor of integers
        logits = self.token_embedding_table(idx)  # (B,T,C)
        return logits

model = BigramLanguageModel(vocab_size)

# 示例输入索引
idx = torch.tensor([
    [0, 1, 2],
    [3, 4, 0]
])  # 形状为 (2, 3)

# 前向传递
logits = model(idx)

print(logits)